import rospy
from sensor_msgs.msg import PointCloud2
import sensor_msgs.point_cloud2 as pc2
import numpy as np
import rospy
import os
import struct
import argparse

# 固定的路径前缀
BASE_PATH = '/home/zhihe/Documents/Dataset/KITTI/point cloud/'

def read_timestamps(times_path):
    with open(times_path, 'r') as f:
        timestamps = []
        for line in f.readlines():
            try:
                # 将秒为单位的浮点数转换为 rospy.Time 对象
                time_in_seconds = float(line.strip())
                timestamps.append(rospy.Time.from_sec(time_in_seconds))
            except ValueError:
                rospy.logerr(f"无法解析时间戳: {line.strip()}")
                continue
    return timestamps

def publish_pointclouds(folder_path, topic, times_path):
    rospy.init_node('pointcloud_publisher', anonymous=True)
    pub = rospy.Publisher(topic, PointCloud2, queue_size=1000)

    # 读取时间戳
    timestamps = read_timestamps(times_path)

    # 获取点云文件列表并排序
    file_list = sorted([f for f in os.listdir(folder_path) if f.endswith('.bin')])

    if len(file_list) != len(timestamps):
        rospy.logerr("点云文件数量与时间戳数量不一致！")
        return

    rate = rospy.Rate(10)  # 发布频率为10Hz
    
    while pub.get_num_connections() == 0 and not rospy.is_shutdown():
        rospy.sleep(0.1)

    for idx, file_name in enumerate(file_list):
        file_path = os.path.join(folder_path, file_name)

        # 读取二进制点云文件
        with open(file_path, 'rb') as f:
            data = f.read()

        # 假设点云数据的字段定义
        fields = [
            pc2.PointField(name='x', offset=0, datatype=pc2.PointField.FLOAT32, count=1),
            pc2.PointField(name='y', offset=4, datatype=pc2.PointField.FLOAT32, count=1),
            pc2.PointField(name='z', offset=8, datatype=pc2.PointField.FLOAT32, count=1),
            pc2.PointField(name='intensity', offset=12, datatype=pc2.PointField.FLOAT32, count=1)
        ]

        # 获取点云数据的点数
        num_points = len(data) // struct.calcsize('ffff')

        # 创建 PointCloud2 消息
        header = rospy.Header()
        header.stamp = timestamps[idx]  # 使用秒单位的时间戳
        header.frame_id = 'map'  # 设置坐标系

        pointcloud_data = np.frombuffer(data, dtype=np.float32)
        pointcloud_data = pointcloud_data.reshape((num_points, 4))

        cloud_msg = pc2.create_cloud(header, fields, pointcloud_data.tolist())

        # 发布点云消息
        pub.publish(cloud_msg)
        rospy.loginfo(f"Published point cloud: {file_path} with timestamp {header.stamp.to_sec()}")

        rate.sleep()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Publish point clouds with corresponding timestamps.")
    parser.add_argument("--folder", type=str, required=True, help="Suffix of the folder (e.g., '08').")
    parser.add_argument("--topic", type=str, default="/velodyne_points", help="ROS topic name for point clouds.")

    args = parser.parse_args()

    # 拼接完整路径
    folder_path = BASE_PATH + args.folder + '/velodyne'
    times_path = BASE_PATH + args.folder + '/times.txt'
    topic = args.topic

    try:
        publish_pointclouds(folder_path, topic, times_path)
    except rospy.ROSInterruptException:
        pass
